# core/aigc_integration.py
import os
import http.client
import json
import requests
from typing import Dict, Any, List
from config.settings import Settings
from utils.aigc_utils import AIGCProgressBar
from utils.file_utils import save_json, load_json
from openai import OpenAI


class CozeAssistant:
    def __init__(self, api_key: str = None, model: str = None):
        self.api_key = api_key or Settings.COZE_API_KEY
        self.model = model or Settings.COZE_MODEL
        self.base_url = Settings.COZE_API_ENDPOINT
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def generate(self, prompt: str, max_tokens: int = 1500, temperature: float = 0.7) -> str:
        """使用Aigc生成内容"""
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature
        }

        try:
            response = requests.post(
                self.base_url,
                headers=self.headers,
                json=payload,
                timeout=Settings.TIMEOUT
            )
            response.raise_for_status()

            # 更健壮的响应处理
            response_data = response.json()

            # 检查不同的响应格式
            if "choices" in response_data and len(response_data["choices"]) > 0:
                return response_data['choices'][0]['message']['content']
            elif "output" in response_data:
                return response_data["output"]
            elif "text" in response_data:
                return response_data["text"]
            else:
                print(f"⚠️ 未知的响应格式: {json.dumps(response_data, indent=2)}")
                return ""

        except requests.exceptions.RequestException as e:
            print(f"❌ Coze API 请求错误: {e}")
            return ""
        except (KeyError, TypeError) as e:
            print(f"❌ Coze API 响应解析错误: {e}")
            return ""
        except json.JSONDecodeError:
            print(f"❌ Coze API 响应不是有效的JSON: {response.text}")
            return ""

    def generate_test_case(self, feature_desc: str) -> Dict[str, Any]:
        """生成测试用例"""
        prompt = f"""
        作为高级QA工程师，请为以下功能创建详细的测试用例:
        {feature_desc}

        返回格式:
        {{
            "test_name": "描述性测试名称",
            "steps": ["步骤1", "步骤2"],
            "expected": "预期结果",
            "priority": "高/中/低",
            "tags": ["标签1", "标签2"],
            "code": "使用pytest实现的完整测试代码"
        }}
        """
        response = self.generate(prompt, max_tokens=2000)
        try:
            return json.loads(response.strip())
        except json.JSONDecodeError:
            print("❌ 解析测试用例JSON失败")
            return {}

    def analyze_failure(self, error_log: str, test_code: str) -> Dict[str, str]:
        """分析测试失败原因"""
        prompt = f"""
        测试失败分析:
        错误日志:
        {error_log}

        测试代码:
        {test_code}

        请分析:
        1. 失败的根本原因
        2. 修复建议
        3. 改进后的测试代码（如果有必要）

        返回JSON格式:
        {{
            "root_cause": "分析原因",
            "fix_suggestion": "修复建议",
            "improved_code": "改进后的代码"
        }}
        """
        response = self.generate(prompt, max_tokens=2000)
        try:
            return json.loads(response.strip())
        except json.JSONDecodeError:
            print("❌ 解析分析结果JSON失败")
            return {}

    def generate_test_data(self, schema: Dict, count: int = 10) -> List[Dict]:
        """生成测试数据"""
        prompt = f"""
        根据以下数据模式生成{count}组多样化的测试数据:
        数据模式: {json.dumps(schema, indent=2)}

        要求:
        - 包含边界值
        - 包含异常值
        - 包含典型值
        - 格式为JSON数组

        返回格式:
        [
            {{"field1": value1, "field2": value2, ...}},
            ...
        ]
        """
        response = self.generate(prompt, max_tokens=2000)
        try:
            return json.loads(response.strip())
        except json.JSONDecodeError:
            print("❌ 解析测试数据JSON失败")
            return []


# 此处生成 DEEPSEEK的 DEEPSEEKAssistant 类
class DEEPSEEKAssistant:
    def __init__(self, api_key: str = None, model: str = None):
        self.api_key = api_key or Settings.DEEPSEEK_API_KEY
        self.model = model or Settings.DEEPSEEK_MODEL
        self.base_url = Settings.DEEPSEEK_API_ENDPOINT  # 应为 "api.deepseek.com"

    def generate(self, prompt: str, max_tokens: int = 1024, temperature: float = 0.7, stream: bool = False) -> str:
        """
        使用 DeepSeek API 生成文本

        Args:
            prompt: 输入提示文本
            max_tokens: 最大输出 token 数量
            temperature: 输出随机性控制参数 (0-1)
            stream: 是否启用流式输出（当前未实现）

        Returns:
            生成的文本内容
        """
        print("DeepSeek AI 请求开始\n")
        # 此处增加一个等待进度条，格式如下：
        # AIGC请求中，请静待返回 [======>      ]  5s
        progress_bar = AIGCProgressBar()
        progress_bar.start()
        # 构建请求体
        payload = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": prompt}
            ],
            "model": self.model,
            "frequency_penalty": 0,
            "max_tokens": max_tokens,
            "presence_penalty": 0,
            "response_format": {"type": "text"},
            "stop": None,
            "stream": stream,
            "stream_options": None,
            "temperature": temperature,
            "top_p": 1,
            "tools": None,
            "tool_choice": "none",
            "logprobs": False,
            "top_logprobs": None
        }

        # 设置请求头
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }

        try:
            # 创建 HTTPS 连接
            conn = http.client.HTTPSConnection(self.base_url)

            # 发送 POST 请求
            conn.request(
                method="POST",
                url="/chat/completions",
                body=json.dumps(payload),
                headers=headers
            )

            # 获取响应
            response = conn.getresponse()

            # 读取响应数据
            data = response.read()

            # 关闭连接
            conn.close()

            # 解析 JSON 响应
            result = json.loads(data.decode("utf-8"))

            progress_bar.stop()
            # 返回生成的文本
            return result["choices"][0]["message"]["content"]

        except Exception as e:
            print(f"API 请求失败: {str(e)}")
            return ""

    def analyze_failure(self, error_log: str, test_code: str) -> Dict[str, str]:
        """分析测试失败原因"""
        prompt = f"""
        测试失败分析:
        错误日志:
        {error_log}

        测试代码:
        {test_code}

        请分析:
        1. 失败的根本原因
        2. 修复建议
        3. 改进后的测试代码（如果有必要）

        返回JSON格式:
        {{
            "root_cause": "分析原因",
            "fix_suggestion": "修复建议",
            "improved_code": "改进后的代码"
        }}
        """
        response = self.generate(prompt, max_tokens=2000)
        try:
            return json.loads(response.strip())
        except json.JSONDecodeError:
            print("❌ 解析分析结果JSON失败")
            return {}

# 此处生成通过阿里云百炼调用 DEEPSEEK的 DEEPSEEKAssistant 类
class DashscopeAssistant:
    def __init__(self, api_key: str = None, model: str = None):
        self.api_key = api_key or Settings.DASHSCOPE_API_KEY
        self.model = model or Settings.DASHSCOPE_MODEL
        self.base_url = Settings.DASHSCOPE_API_ENDPOINT

    def generate(self, prompt: str, max_tokens: int = 1024, temperature: float = 0.7, stream: bool = False) -> str:
        """
        使用 DASHSCOPE DeepSeek API 生成文本

        Args:
            prompt: 输入提示文本
            max_tokens: 最大输出 token 数量
            temperature: 输出随机性控制参数 (0-1)
            stream: 是否启用流式输出（当前未实现）

        Returns:
            生成的文本内容
        """
        print("DASHSCOPE DeepSeek AI 请求开始\n")
        # 此处增加一个等待进度条，格式如下：
        # AIGC请求中，请静待返回 [======>      ]  5s
        progress_bar = AIGCProgressBar()
        progress_bar.start()
        # 构建请求体
        payload = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": prompt}
            ],
            "model": self.model,
            "frequency_penalty": 0,
            "max_tokens": max_tokens,
            "presence_penalty": 0,
            "response_format": {"type": "text"},
            "stop": None,
            "stream": stream,
            # "stream_options": None,
            "temperature": temperature,
            "top_p": 1,
            # "tools": None,
            # "tool_choice": "none",
            # "logprobs": False,
            # "top_logprobs": None
        }

        # 设置请求头
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }

        try:
            # 创建 HTTPS 连接
            conn = http.client.HTTPSConnection(self.base_url)
            # 发送 POST 请求
            conn.request(
                method="POST",
                url="/compatible-mode/v1/chat/completions",
                body=json.dumps(payload),
                headers=headers
            )

            # 获取响应
            response = conn.getresponse()

            # 读取响应数据
            data = response.read()

            # 关闭连接
            conn.close()

            # 解析 JSON 响应
            result = json.loads(data.decode("utf-8"))

            progress_bar.stop()
            # 返回生成的文本
            return result["choices"][0]["message"]["content"]

        except Exception as e:
            print(f"API 请求失败: {str(e)}")
            return ""

    def analyze_failure(self, error_log: str, test_code: str) -> Dict[str, str]:
        """分析测试失败原因"""
        prompt = f"""
        测试失败分析:
        错误日志:
        {error_log}

        测试代码:
        {test_code}

        请分析:
        1. 失败的根本原因
        2. 修复建议
        3. 改进后的测试代码（如果有必要）

        返回JSON格式:
        {{
            "root_cause": "分析原因",
            "fix_suggestion": "修复建议",
            "improved_code": "改进后的代码"
        }}
        """
        response = self.generate(prompt, max_tokens=2000)
        try:
            return json.loads(response.strip())
        except json.JSONDecodeError:
            print("❌ 解析分析结果JSON失败")
            return {}